// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/sparksql/aggregates/register.h>

using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::functions::aggregate::sparksql::test {

namespace {

class CollectListAggregateTest : public AggregationTestBase {
 protected:
  void SetUp() override {
    AggregationTestBase::SetUp();
    registerAggregateFunctions("spark_");
  }
};

TEST_F(CollectListAggregateTest, groupBy) {
  std::vector<RowVectorPtr> batches;
  // Creating 3 batches of input data.
  // 0: {0, null} {0, 1}    {0, 2}
  // 1: {1, 1}    {1, null} {1, 3}
  // 2: {2, 2}    {2, 3}    {2, null}
  // 3: {3, 3}    {3, 4}    {3, 5}
  // 4: {4, 4}    {4, 5}    {4, 6}
  for (auto i = 0; i < 3; i++) {
    RowVectorPtr data = make_row_vector(
        {make_flat_vector<int32_t>({0, 1, 2, 3, 4}),
         make_flat_vector<int64_t>(
             5,
             [&i](const vector_size_t& row) { return i + row; },
             [&i](const auto& row) { return i == row; })});
    batches.push_back(data);
  }

  auto expected = make_row_vector(
      {make_flat_vector<int32_t>({0, 1, 2, 3, 4}),
       make_array_vector_from_json<int64_t>(
           {"[1, 2]", "[1, 3]", "[2, 3]", "[3, 4, 5]", "[4, 5, 6]"})});

  testAggregations(
      batches,
      {"c0"},
      {"spark_collect_list(c1)"},
      {"c0", "array_sort(a0)"},
      {expected});
  testAggregationsWithCompanion(
      batches,
      [](auto& /*builder*/) {},
      {"c0"},
      {"spark_collect_list(c1)"},
      {{BIGINT()}},
      {"c0", "array_sort(a0)"},
      {expected},
      {});
}

TEST_F(CollectListAggregateTest, global) {
  auto data = make_row_vector({make_nullable_flat_vector<int32_t>(
      {std::nullopt, 1, 2, std::nullopt, 4, 5})});
  auto expected =
      make_row_vector({make_array_vector_from_json<int32_t>({"[1, 2, 4, 5]"})});

  testAggregations(
      {data}, {}, {"spark_collect_list(c0)"}, {"array_sort(a0)"}, {expected});
  testAggregationsWithCompanion(
      {data},
      [](auto& /*builder*/) {},
      {},
      {"spark_collect_list(c0)"},
      {{INTEGER()}},
      {"array_sort(a0)"},
      {expected});
}

TEST_F(CollectListAggregateTest, ignoreNulls) {
  auto input = make_row_vector({make_nullable_flat_vector<int32_t>(
      {1, 2, std::nullopt, 4, std::nullopt, 6})});
  // Spark will ignore all null values in the input.
  auto expected =
      make_row_vector({make_array_vector_from_json<int32_t>({"[1, 2, 4, 6]"})});
  testAggregations(
      {input}, {}, {"spark_collect_list(c0)"}, {"array_sort(a0)"}, {expected});
  testAggregationsWithCompanion(
      {input},
      [](auto& /*builder*/) {},
      {},
      {"spark_collect_list(c0)"},
      {{INTEGER()}},
      {"array_sort(a0)"},
      {expected},
      {});
}

TEST_F(CollectListAggregateTest, allNullsInput) {
  auto input = make_row_vector({make_all_null_flat_vector<int64_t>(100)});
  // If all input data is null, Spark will output an empty array.
  auto expected = make_row_vector({make_array_vector_from_json<int32_t>({"[]"})});
  testAggregations({input}, {}, {"spark_collect_list(c0)"}, {expected});
  testAggregationsWithCompanion(
      {input},
      [](auto& /*builder*/) {},
      {},
      {"spark_collect_list(c0)"},
      {{BIGINT()}},
      {},
      {expected},
      {});
}
} // namespace
} // namespace kumo::pollux::functions::aggregate::sparksql::test
